{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Part 4, Topic 2: CPA on Firmware Implementation of AES"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**SUMMARY**: *By now, you'll have used a DPA attack to break AES. While this method has its place in side channel attacks, it often requires a large number of traces to break AES and can suffer from additional issues like ghost peaks.*\n",
    "\n",
    "*We've also learned in the previous lab that there is a very linear relationship between the hamming weight of the SBox output and the power consumption at that point. Instead of checking average power consumption over many traces to see if a guessed subkey is correct, we can instead check if our guessed subkey also has this linear relationship with the device's power consumption across a set of traces. Like with DPA, we'll need to repeat this measurement at each point in time along the power trace.*\n",
    "\n",
    "*To get an objective measurement of how linear this relationship is, we'll be developing some code to calculate the Pearson correlation coefficient.*\n",
    "\n",
    "**LEARNING OUTCOMES:**\n",
    "* Developing an algorithm based on a mathematical description\n",
    "* Verify that correlation can be used to break a single byte of AES\n",
    "* Extend the single byte attack to the rest of the key"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Prerequisites"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This notebook will build upon previous ones. Make sure you've completed the following tutorials and their prerequisites:\n",
    "\n",
    "* ☑ Part 3 notebooks (you should be comfortable with running an attack on AES)\n",
    "* ☑ Power and Hamming Weight Relationship (we'll be using information from this tutorial)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## AES Trace Capture"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Our first step will be to send some plaintext to the target device and observe its power consumption during the encryption. The capture loop will be the same as in the DPA attack. This time, however, we'll only need 50 traces to recover the key, a major improvement over the last attack!\n",
    "\n",
    "Depending what you are using, you can complete this either by:\n",
    "\n",
    "* Capturing new traces from a physical device.\n",
    "* Reading pre-recorded data from a file.\n",
    "\n",
    "You get to choose your adventure - see the two notebooks with the same name of this, but called `(SIMULATED)` or `(HARDWARE)` to continue. Inside those notebooks you should get some code to copy into the following section, which will define the capture function.\n",
    "\n",
    "Be sure you get the `\"✔️ OK to continue!\"` print once you run the cell afterwards, otherwise things will fail later on!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# There is a file with the same name of this lab but (HARDWARE) in title for using CW-Nano/CW-Lite/CW-Pro\n",
    "# There is a file with the same name of this lab but (SIMULATED) in title for using recorded data\n",
    "raise NotImplementedError(\"Insert code from (HARDWARE) or (SIMULATED) Notebook Here\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "assert len(trace_array) == 50\n",
    "print(\"✔️ OK to continue!\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Again, let's quickly plot a trace to make sure everything looks as expected:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib notebook\n",
    "import matplotlib.pylab as plt\n",
    "\n",
    "# ###################\n",
    "# Add your code here\n",
    "# ###################\n",
    "raise NotImplementedError(\"Add your code here, and delete this.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## AES Model and Hamming Weight"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Like with the previous tutorial, we'll need to be able to easily grab what the sbox output will be for a given plaintext and key, as well as get the hamming weight of numbers between 0 and 255:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ###################\n",
    "# Add your code here\n",
    "# ###################\n",
    "raise NotImplementedError(\"Add your code here, and delete this.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Verify that your model is correct:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "assert HW[aes_internal(0xA1, 0x79)] == 3\n",
    "assert HW[aes_internal(0x22, 0xB1)] == 5\n",
    "print(\"✔️ OK to continue!\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Developing our Correlation Algorithm "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As we discussed earlier, we'll be testing how good our guess is using a measurement called the Pearson correlation coefficient, which measures the linear correlation between two datasets. \n",
    "\n",
    "The actual algorithm is as follows for datasets $X$ and $Y$ of length $N$, with means of $\\bar{X}$ and $\\bar{Y}$, respectively:\n",
    "\n",
    "$$r = \\frac{cov(X, Y)}{\\sigma_X \\sigma_Y}$$\n",
    "\n",
    "$cov(X, Y)$ is the covariance of `X` and `Y` and can be calculated as follows:\n",
    "\n",
    "$$cov(X, Y) = \\sum_{n=1}^{N}[(Y_n - \\bar{Y})(X_n - \\bar{X})]$$\n",
    "\n",
    "$\\sigma_X$ and $\\sigma_Y$ are the standard deviation of the two datasets. This value can be calculated with the following equation:\n",
    "\n",
    "$$\\sigma_X = \\sqrt{\\sum_{n=1}^{N}(X_n - \\bar{X})^2}$$\n",
    "\n",
    "As you can see, the calulation is actually broken down pretty nicely into some smaller chunks that we can implement with some simple functions. While we could use a library to calculate all this stuff for us, being able to implement a mathematical algorithm in code is a useful skill to develop. \n",
    "\n",
    "To start, build the following functions:\n",
    "\n",
    "1. `mean(X)` to calculate the mean of a dataset (the mean being `X_bar` that will be used elsewhere).\n",
    "1. `std_dev(X, X_bar)` to calculate the standard deviation of a dataset. We'll need to reuse the mean for the covariance, so it makes more sense to calculate it once and pass it in to each function\n",
    "1. `cov(X, X_bar, Y, Y_bar)` to calculate the covariance of two datasets. Again, we can just pass in the means we calculate for std_dev here.\n",
    "\n",
    "**HINT: You can use `np.sum(X, axis=0)` to replace all of the $\\sum$ from earlier. The argument `axis=0` will sum across columns, allowing us to use a single `mean`, `std_dev`, and `cov` call for the entire power trace**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def mean(X):\n",
    "    raise NotImplementedError(\"Add your code here, and delete this.\")\n",
    "\n",
    "def std_dev(X, X_bar):\n",
    "    raise NotImplementedError(\"Add your code here, and delete this.\")\n",
    "\n",
    "def cov(X, X_bar, Y, Y_bar):\n",
    "    raise NotImplementedError(\"Add your code here, and delete this.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's quickly check to make sure everything's as expected. The following blocks will run some test vectors on your functions, confirm you get the correct answer:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "a = np.array([[5, 3, 4, 4, 5, 6],\n",
    "             [27, 2, 3, 4, 12, 6],\n",
    "              [1, 3, 5, 4, 5, 6],\n",
    "              [1, 2, 3, 4, 5, 6],\n",
    "             ]).transpose()\n",
    "a_bar = mean(a)\n",
    "b = np.array([[5, 4, 3, 2, 1, 3]]).transpose()\n",
    "b_bar = mean(b)\n",
    "\n",
    "o_a = std_dev(a, a_bar)\n",
    "o_b = std_dev(b, b_bar)\n",
    "\n",
    "ab_cov = cov(a, a_bar, b, b_bar)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "assert (a_bar == np.array([4.5, 9., 4., 3.5])).all()\n",
    "assert (b_bar == np.array([3.])).all()\n",
    "assert (o_a[3] > 4.1833001 and o_a[3] < 4.1833002)\n",
    "assert (o_b[0] > 3.162277 and o_b[0] < 3.162278)\n",
    "assert (ab_cov == np.array([-1., 28., -9., -10.])).all()\n",
    "print(\"✔️ OK to continue!\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now that we've got all the building blocks to our correlation function, let's see if we can put everything together and break a single byte of AES. In order to do this, let's take a closer look at what we're trying to do and the data we've got."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Correlation Data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Remember that the general correlation formula for two datasets $X$  and $Y$ is:\n",
    "\n",
    "$$r = \\frac{cov(X, Y)}{\\sigma_X \\sigma_Y}$$\n",
    "\n",
    "We are going to be correlateing between a power measurment (`trace_array`) and Hamming weight of a key guess. First let's look at our power trace array:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(trace_array)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You should have something like the following:\n",
    "```python\n",
    "[\n",
    "    [point_0, point_1, point_2, ...], # trace 0\n",
    "    [point_0, point_1, point_2, ...], # trace 1\n",
    "    [point_0, point_1, point_2, ...], # trace 2\n",
    "    ...\n",
    "]\n",
    "```\n",
    "\n",
    "where the rows of the array are the different traces we captured and the columns of the array are the different points in those traces. The columns here will be one of the two datasets for our correlation equation. The other dataset will be the hamming weight of the SBox output, for a given *key guess* `key` of the byte we are looking at:\n",
    "\n",
    "```python\n",
    "[\n",
    "      [HW[aes_internal(plaintext0[0], key[0])], # trace 0\n",
    "      [HW[aes_internal(plaintext1[0], key[0])], # trace 1\n",
    "      [HW[aes_internal(plaintext2[0], key[0])], # trace 2\n",
    "      ...\n",
    "]\n",
    "```\n",
    "\n",
    "which we'll shorten to:\n",
    "\n",
    "```python\n",
    "[\n",
    "      [hw], # trace 1\n",
    "      [hw], # trace 2\n",
    "      [hw], # trace 3\n",
    "      ...\n",
    "]\n",
    "```\n",
    "\n",
    "Like with the DPA attack, we don't know where the encryption is occurring, meaning we have to repeat the correlation calculation for each column in the trace array, with the largest correlation being our best guess for where the SBox output is happening. We obviously also don't know the key (that's the thing we're trying to find!), so we'll also need to repeat the best correlation calculation for each possible value of `key[0]` (0 to 255). The key with the highest absolute correlation is our best guess for the value of the key byte."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Correlation Attack Implementaiton"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The correlation attack is basically to calculate this:\n",
    "\n",
    "$$r = \\frac{cov(X, Y)}{\\sigma_X \\sigma_Y}$$\n",
    "\n",
    "Where:\n",
    "\n",
    "* $X$ is a power trace sample point\n",
    "* $Y$ is an internal state guess\n",
    "\n",
    "Remember you already defined (and tested) the functions that generate `cov(X,Y)`, and also the `std_dev(X)` ($\\sigma_X$). The actual API for those functions requires to to pass in the `mean()` as well as a seperate argument (passed in for computational efficiency, since it is re-used).\n",
    "\n",
    "### Hint: Using Vectors\n",
    "\n",
    "We should mention a few way to improve your work.\n",
    "\n",
    "A really nice feature of numpy is that we can do the correlation calculations across the entire trace at once (mean, std_dev, cov). That means there's no need to do:\n",
    "\n",
    "```python\n",
    "t_bar = []\n",
    "for point_num in range(len(trace_array[0])):\n",
    "    t_bar.append(mean(trace_array[:,point_num]))\n",
    "    # and so on...\n",
    "\n",
    "t_bar = np.array(t_bar)\n",
    "```\n",
    "\n",
    "when we can do\n",
    "\n",
    "```python\n",
    "t_bar = mean(trace_array)\n",
    "```\n",
    "\n",
    "and get the same thing back. The only caveat being that we need to make sure that the columns and rows of our arrays are the right way around (i.e. make sure your hamming weight array has 1 column and 50 rows and not the other way around). If you find it easier to construct and array one way and not the other, you can use the `.transpose()` method to swap the rows and columns.\n",
    "\n",
    "### Finding Largest Correlation\n",
    "\n",
    "Once you've got all your correlations for a particular key guess, you want to find the largest absolute correlation. We're taking the absolute value of the correlation here since we only care that the relation between hamming weight and the power trace is linear, not that the slope is positive or negative. `max(abs(correlations))` will do that for you.\n",
    "\n",
    "### Enumerating Guesses\n",
    "\n",
    "Perform this for every possible value of the key byte (aka 0 to 255) and the one with the largest correlation is your best guess for the key. It's up to you how you want to extract this information from your loop, but one way of doing it is to stick the best guess for each of your key guesses in an array. Once you've gone through all the key guesses, you can extract the best guess with `np.argmax(maxcpa)` and the correlation of that guess with `max(maxcpa)`.\n",
    "\n",
    "### Equation to Python\n",
    "\n",
    "We can take the earlier equation and plug in some of our Python variable names to give you a good starting point. We are using:\n",
    "\n",
    "* $r$ = `cpaoutput`\n",
    "* $X$ = `t` or `trace_array` (the average of it called `t_bar`).\n",
    "* $Y$ = `hws` (the mean of it called `hws_bar`).\n",
    "\n",
    "Our equation now looks something like this:\n",
    "\n",
    "$$cpaoutput = \\frac{cov(X, Y)}{\\sigma_X \\sigma_Y}$$\n",
    "\n",
    "This should almost directly convert to Python code!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "maxcpa = [0] * 256\n",
    "\n",
    "# we don't need to redo the mean and std dev calculations \n",
    "# for each key guess\n",
    "t_bar = mean(trace_array) \n",
    "o_t = std_dev(trace_array, t_bar)\n",
    "\n",
    "for kguess in trange(0, 256):\n",
    "    hws = np.array([[HW[aes_internal(textin[0],kguess)] for textin in textin_array]]).transpose()\n",
    "    \n",
    "    # ###################\n",
    "    # Add your code here  \n",
    "    # ###################\n",
    "    cpaoutput = ???\n",
    "    maxcpa[kguess] = ???  \n",
    "    raise NotImplementedError(\"Add your code here, and delete this.\")\n",
    "    \n",
    "print(\"Key guess: \", hex(guess))\n",
    "print(\"Correlation: \", guess_corr)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's make sure we've recovered the byte correctly:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "assert guess == 0x2b\n",
    "print(\"✔️ OK to continue!\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To break the rest of the key, simply repeat the attack for the rest of the bytes of the key. Don't forget to update your code from above to use the correct byte of the plaintext!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "t_bar = np.sum(trace_array, axis=0)/len(trace_array)\n",
    "o_t = np.sqrt(np.sum((trace_array - t_bar)**2, axis=0))\n",
    "\n",
    "cparefs = [0] * 16 #put your key byte guess correlations here\n",
    "bestguess = [0] * 16 #put your key byte guesses here\n",
    "\n",
    "for bnum in trange(0, 16):\n",
    "    maxcpa = [0] * 256\n",
    "    for kguess in range(0, 256):\n",
    "        # ###################\n",
    "        # Add your code here\n",
    "        # ###################\n",
    "        raise NotImplementedError(\"Add your code here, and delete this.\")\n",
    "\n",
    "print(\"Best Key Guess: \", end=\"\")\n",
    "for b in bestguess: print(\"%02x \" % b, end=\"\")\n",
    "print(\"\\n\", cparefs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "With one final check to make sure you've got the correct key:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for bnum in range(16):\n",
    "    assert bestguess[bnum] == key[bnum], \\\n",
    "    \"Byte {} failed, expected {:02X} got {:02X}\".format(bnum, key[bnum], bestguess[bnum])\n",
    "print(\"✔️ OK to continue!\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We're done! There's actually a lot of room to expand on this attack:\n",
    "\n",
    "1. Currently, the loop needs to go through all the traces before it can return a correlation. This isn't too bad for a short attack, for a much longer one (think 10k+ traces) we won't get any feedback from the attack until it's finished. Also, if we didn't capture enough traces for the attack, the entire analysis calculation needs to be repeated! Instead of using the original correlation equation, we can instead use an equivalent \"online\" version that can be easily updated with more traces: $$r_{i,j} = \\frac{D\\sum_{d=1}^{D}h_{d,i}t_{d,j}-\\sum_{d=1}^{D}h_{d,i}\\sum_{d=1}^{D}t_{d,j}}{\\sqrt{((\\sum_{d=1}^Dh_{d,i})^2-D\\sum_{d=1}^Dh_{d,i}^2)-((\\sum_{d=1}^Dt_{d,j})^2-D\\sum_{d=1}^Dh_{d,j}^2)}}$$\n",
    "where\n",
    "\n",
    "| **Equation** | **Python Variable** | **Value**  | \n",
    "|--------------|---------------------|------------|\n",
    "|  d           |       tnum          | trace number |\n",
    "|  i           |       kguess        | subkey guess |\n",
    "| j | j index trace point | sample point in trace |\n",
    "| h | hypint | guess for power consumption | \n",
    "| t | traces | traces | \n",
    "\n",
    "2. There's a lot more we can learn from the attack other than the key. For example, we could plot how far away the correct key guess is from the top spot (called the partial guessing entropy or PGE) vs. how many traces we used, giving us a better idea of how many traces we needed to actually recover the correct key. We also might want to plot how correlation for a given key guess changes over time.\n",
    "\n",
    "This \"online\" correlation equation is the one that the subject of the next tutorial, ChipWhisperer Analyzer, actually uses. It also provides functions and methods for gathering and plotting some interesting statistics."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "<small>NO-FUN DISCLAIMER: This material is Copyright (C) NewAE Technology Inc., 2015-2020. ChipWhisperer is a trademark of NewAE Technology Inc., claimed in all jurisdictions, and registered in at least the United States of America, European Union, and Peoples Republic of China.\n",
    "\n",
    "Tutorials derived from our open-source work must be released under the associated open-source license, and notice of the source must be *clearly displayed*. Only original copyright holders may license or authorize other distribution - while NewAE Technology Inc. holds the copyright for many tutorials, the github repository includes community contributions which we cannot license under special terms and **must** be maintained as an open-source release. Please contact us for special permissions (where possible).\n",
    "\n",
    "THE SOFTWARE IS PROVIDED \"AS IS\", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.</small>"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
